import pandas as pd
import lightgbm as lgb
import joblib
from sklearn.metrics import mean_squared_error

from . import config
from . import data_loader
from . import feature_engineering
from . import evaluation

def train_model():
    """
    完整的模型训练流程：加载数据 -> 特征工程 -> 训练 -> 评估 -> 保存
    """
    # 1. 加载数据
    df = data_loader.load_all_net_flows_from_redis()
    if df is None:
        return

    # 2. 特征工程
    df = feature_engineering.create_time_features(df)
    df = feature_engineering.create_lag_features(df, config.LAG_WINDOW_SIZE)

    # 3. 数据清洗和准备
    df.dropna(inplace=True)
    
    # --- 核心修改：划分训练集和测试集 ---
    # 我们使用最后两天的数据作为测试集，模拟真实预测场景
    split_date = df['timestamp'].max() - pd.Timedelta(days=2)
    train_df = df[df['timestamp'] < split_date]
    test_df = df[df['timestamp'] >= split_date]

    print(f"\n数据划分完成:")
    print(f"  - 训练集大小: {len(train_df)} 条")
    print(f"  - 测试集大小: {len(test_df)} 条")

    # 定义特征 (X) 和目标 (y)
    y_train = train_df['net_flow']
    X_train = train_df.drop(columns=['timestamp', 'net_flow'])
    
    y_test = test_df['net_flow']
    X_test = test_df.drop(columns=['timestamp', 'net_flow'])
    
    categorical_features = ['community_id', 'hour', 'day_of_week', 'is_weekend', 'time_slot_of_day']

    print("\n训练集特征预览:")
    print(X_train.head(5))

    print("\n数据准备完成，开始训练模型...")
    
    # 4. 训练模型
    model = lgb.LGBMRegressor(
        objective='regression_l1',
        n_estimators=1000,
        learning_rate=0.05,
        num_leaves=31,
        n_jobs=-1,
        seed=42
    )

    # 模型在训练时会监控在测试集上的表现，如果连续100轮没有提升，就提前停止训练
    model.fit(
        X_train, y_train,
        eval_set=[(X_test, y_test)],
        eval_metric='rmse',
        categorical_feature=categorical_features,
        callbacks=[lgb.early_stopping(100, verbose=True)]
    )

    print("\n✅ 模型训练完成！")

    # 5. 在测试集上进行预测和评估
    predictions = model.predict(X_test)
    
    # 将预测结果添加到测试集的 DataFrame 中，方便后续处理
    test_df = test_df.copy() # 避免 SettingWithCopyWarning
    test_df['prediction'] = predictions
    
    # 调用评估函数，打印性能指标
    evaluation.evaluate_model_performance(y_test, predictions)
    
    # 调用可视化函数，生成对比图
    evaluation.plot_predictions_vs_actuals(test_df, community_id_to_plot=1)

    # 6. 保存模型
    try:
        joblib.dump(model, config.MODEL_OUTPUT_PATH)
        print(f"\n模型已成功保存到: {config.MODEL_OUTPUT_PATH}")
    except Exception as e:
        print(f"❌ 模型保存失败: {e}")

if __name__ == "__main__":
    train_model()